{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cf29e480d77ef507",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "\n",
    "# 使用stream\n",
    "流式处理对于基于 LLM 的应用程序对最终用户的响应至关重要。\n",
    "\n",
    "重要的 LangChain 原语，如 LLMs、解析器、提示、检索器和代理实现了 LangChain [Runnable 接口](https://python.langchain.com.cn/docs/expression_language/interface)。\n",
    "\n",
    "该接口提供了两种常见的流式内容的方法：\n",
    "\n",
    "1. `sync stream` 和 `asyn astream`：流式处理的**默认实现**，从链中流式传输**最终输出**。\n",
    "2. `async astream_events` 和 `async astream_log`：这些方法提供了一种从链中流式传输**中间步骤**和**最终输出**的方式。\n",
    "\n",
    "\n",
    "## LLMs 和 Chat Models\n",
    "大型语言模型生成对查询的完整响应可能需要几秒钟。这远远慢于应用程序对最终用户响应感觉灵敏的~200-300 ms的阈值。\n",
    "使应用程序感觉更灵敏的关键策略是显示中间进度；例如，逐个令牌从模型中流式传输输出。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fd5a9e7d8a49a2f5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T08:26:17.734435Z",
     "start_time": "2024-10-25T08:26:10.807003Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|你好！|我是|Q|wen，一个由|阿里云开发的|大型语言模型。|我的设计目的是为了|帮助用户生成各种|类型的文本，如|文章、故事、|诗歌、故事等|，并能够根据用户的|指示完成诸如撰写|邮件、表达观点|、玩游戏等多种任务|。我通过学习|大量的文本数据来|理解人类的语言和|知识体系，但|请注意，我的知识|可能具有一定的时效|限制，对于一些|最新发生的事情或|非常具体的专业知识|，我可能无法|提供最准确的信息|。此外，虽然|我可以处理多种语言|，但我主要被|训练使用中文进行|交流。希望我能|成为你工作和|生活中的好帮|手！如果你有任何|问题想要询问或者|需要帮助的地方，请|随时告诉我。||"
     ]
    }
   ],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "import os\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "llm = ChatOpenAI(\n",
    "    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key=\"sk-xxx\",\n",
    "    openai_api_key=os.getenv(\"DASHSCOPE_API_KEY\"),\n",
    "    openai_api_base=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n",
    "    model_name=\"qwen-max\"\n",
    ")\n",
    "chunks = []\n",
    "async for chunk in llm.astream(\"你好。告诉我一些关于你自己的事情\"):\n",
    "    chunks.append(chunk)\n",
    "    print(chunk.content, end=\"|\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0c848e24e5851b9",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 链\n",
    "几乎所有的 LLM 应用都涉及到不止一个调用语言模型的步骤。\n",
    "让我们使用 LangChain 表达语言 (LCEL) 创建一个简单的链，它结合了一个提示、模型和一个解析器，并验证了流式处理是否有效。\n",
    "我们将使用 StrOutputParser 来解析模型的输出。这是一个简单的解析器，从 AIMessageChunk 中提取 content 字段，给我们模型返回的 token。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5510bb5e6485ad4d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T08:25:22.887342Z",
     "start_time": "2024-10-25T08:25:17.858533Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|当然，|这里|有一个|关于猫的笑话|供您一笑：\n",
      "\n",
      "|为什么猫咪不喜欢在线|购物？\n",
      "因为它们|更喜欢真正的“|窗”购物！|每当看到窗外的小|鸟或者蝴蝶，|那就是它们最感兴趣的|“商品”啦|！\n",
      "\n",
      "希望这个笑话|能让您会心|一笑。猫咪总是|以它们独特的方式|给我们的生活带来|欢乐。||"
     ]
    }
   ],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "import os\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(\"告诉我一个关于 {topic} 的笑话(200字以内)\")\n",
    "out_parser = StrOutputParser()\n",
    "\n",
    "llm = ChatOpenAI(\n",
    "    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key=\"sk-xxx\",\n",
    "    openai_api_key=os.getenv(\"DASHSCOPE_API_KEY\"),\n",
    "    openai_api_base=\"https://dashscope.aliyuncs.com/compatible-mode/v1\",\n",
    "    model_name=\"qwen-max\"\n",
    ")\n",
    "\n",
    "chain = prompt | llm | out_parser\n",
    "\n",
    "async for chunk in chain.astream({\"topic\": \"猫\"}):\n",
    "    print(chunk, end=\"|\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "345238c0adc40b67",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 使用输入流\n",
    "如果您想要在生成时从输出中流式传输 JSON 怎么办？有一种方法可以做到 -- 解析器需要操作输入流，并尝试将部分 json “自动完成”为有效状态。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ffe446e61c58eb97",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T08:33:36.476345Z",
     "start_time": "2024-10-25T08:33:26.939509Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{}\n",
      "{'joke': {}}\n",
      "{'joke': {'setup': ''}}\n",
      "{'joke': {'setup': '为什么程序员喜欢猫'}}\n",
      "{'joke': {'setup': '为什么程序员喜欢猫？'}}\n",
      "{'joke': {'setup': '为什么程序员喜欢猫？', 'punchline': '因为他们喜欢用'}}\n",
      "{'joke': {'setup': '为什么程序员喜欢猫？', 'punchline': \"因为他们喜欢用'爪'（\"}}\n",
      "{'joke': {'setup': '为什么程序员喜欢猫？', 'punchline': \"因为他们喜欢用'爪'（抓）数据！\"}}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from langchain_core.output_parsers import JsonOutputParser\n",
    "\n",
    "chain = prompt | llm | JsonOutputParser()\n",
    "async for text in chain.astream({\"topic\": \"以JSON格式输出猫笑话\"}):\n",
    "    print(text, flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f658f1f676837427",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 不操作输入流且会中断流式传输的函数\n",
    "现在，让我们中断流式传输。我们将使用先前的示例，并在末尾附加一个提取函数，该函数从最终的 JSON 中提取国家名称。\n",
    ":::{.callout-warning} 链中的任何步骤，如果操作的是最终输入而不是输入流，都可能通过 stream 或 astream 打破流式传输功能。 :::\n",
    ":::{.callout-tip} 稍后，我们将讨论 astream_events API，该 API 将流式传输中间步骤的结果。即使链中包含仅操作最终输入而不是输入流的步骤，该 API 也会流式传输结果。 :::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f2e1dd9dc870fabe",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T08:43:59.299906Z",
     "start_time": "2024-10-25T08:43:45.472560Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['France', 'Spain', 'Japan']|"
     ]
    }
   ],
   "source": [
    "def _extract_country_names(inputs) -> [str]:\n",
    "    \"\"\"\n",
    "    一个不操作输入流并且会中断流式传输的函数。\n",
    "    :param inputs: \n",
    "    :return: \n",
    "    \"\"\"\n",
    "    if not isinstance(inputs, dict):\n",
    "        return \"\"\n",
    "\n",
    "    if \"countries\" not in inputs:\n",
    "        return \"\"\n",
    "\n",
    "    countries = inputs[\"countries\"]\n",
    "\n",
    "    if not isinstance(countries, list):\n",
    "        return \"\"\n",
    "\n",
    "    country_names = [\n",
    "        country.get(\"name\") for country in countries if isinstance(country, dict)\n",
    "    ]\n",
    "\n",
    "    return country_names\n",
    "\n",
    "\n",
    "chain = llm | JsonOutputParser() | _extract_country_names\n",
    "async for text in chain.astream(\n",
    "        '以 JSON 格式输出法国、西班牙和日本的国家及其人口的列表。使用一个带有“countries”外键的字典，其中包含一个国家列表。每个国家应该有“name”和“population”关键字。'\n",
    "):\n",
    "    print(text, end=\"|\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed946a17008dea78",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 生成器函数\n",
    "让我们使用可以操作输入流的生成器函数来修复流式处理。\n",
    "\n",
    ":::{.callout-tip} 生成器函数（使用 yield 的函数）允许编写能够操作输入流的代码。 :::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "80cf58de367dfaa0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T09:13:35.044447Z",
     "start_time": "2024-10-25T09:13:19.476889Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "法国|西班牙|日本|"
     ]
    }
   ],
   "source": [
    "async def _extract_country_names_streaming(input_stream):\n",
    "    \"\"\"\n",
    "    一个操作输入流的生成器函数。\n",
    "    :param input_stream: \n",
    "    :return: \n",
    "    \"\"\"\n",
    "    country_names_so_far = set()\n",
    "\n",
    "    async for input_chunk in input_stream:\n",
    "        if not isinstance(input_chunk, dict):\n",
    "            continue\n",
    "\n",
    "        if \"countries\" not in input_chunk:\n",
    "            continue\n",
    "\n",
    "        countries = input_chunk[\"countries\"]\n",
    "\n",
    "        if not isinstance(countries, list):\n",
    "            continue\n",
    "\n",
    "        for country in countries:\n",
    "            name = country.get(\"name\")\n",
    "            if not name:\n",
    "                continue\n",
    "            if name not in country_names_so_far:\n",
    "                yield name\n",
    "                country_names_so_far.add(name)\n",
    "\n",
    "\n",
    "chain = llm | JsonOutputParser() | _extract_country_names_streaming\n",
    "\n",
    "async for text in chain.astream(\n",
    "        '以 JSON 格式输出法国、西班牙和日本的国家及其人口的列表。使用一个带有“countries”外键的字典，其中包含一个国家列表。每个国家应该有“name”和“population”关键字。'\n",
    "):\n",
    "    print(text, end=\"|\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38b473aa94bb1c4d",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# 使用流事件\n",
    "事件流是一个beta API。该 API可能会根据反馈做出一些更改。\n",
    ":::{.callout-note} 引入于 langchain-core 0.1.14 版本。 :::\n",
    "为了让 `astream_events` API 正常工作：\n",
    "\n",
    "- 尽可能在代码中使用 `async`（例如，异步工具等）\n",
    "- 如果定义自定义函数/可运行对象，请传播回调\n",
    "- 每当使用非 LCEL 的可运行对象时，请确保在 LLM 上调用 `.astream()` 而不是 `.ainvoke`，以强制 LLM 流式传输令牌。\n",
    "- 如果有任何不符合预期的情况，请告诉我们！ :)\n",
    "\n",
    "> 记住，有些组件不适合流式处理，因为它们不适用于**输入流**。\n",
    "虽然这些组件在使用`astream`时可能会中断最终输出的流式处理，但使用`astream_events`仍然会从支持流式处理的中间步骤中产生流式事件！\n",
    "\n",
    "\n",
    "## 聊天模型\n",
    "让我们首先查看聊天模型产生的事件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6156c561e2fb2acc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T09:32:28.177015Z",
     "start_time": "2024-10-25T09:32:27.096408Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'event': 'on_chain_start',\n",
       "  'data': {'input': 'hello'},\n",
       "  'name': 'RunnableSequence',\n",
       "  'tags': [],\n",
       "  'run_id': '1a997380-c991-4732-bb73-7540afcdf5c2',\n",
       "  'metadata': {},\n",
       "  'parent_ids': []},\n",
       " {'event': 'on_chat_model_start',\n",
       "  'data': {'input': {'messages': [[HumanMessage(content='hello', additional_kwargs={}, response_metadata={})]]}},\n",
       "  'name': 'ChatOpenAI',\n",
       "  'tags': ['seq:step:1'],\n",
       "  'run_id': 'b75d8338-57b6-43b2-8f88-a81389844523',\n",
       "  'metadata': {'ls_provider': 'openai',\n",
       "   'ls_model_name': 'qwen-max',\n",
       "   'ls_model_type': 'chat',\n",
       "   'ls_temperature': 0.7},\n",
       "  'parent_ids': ['1a997380-c991-4732-bb73-7540afcdf5c2']},\n",
       " {'event': 'on_chat_model_stream',\n",
       "  'data': {'chunk': AIMessageChunk(content='', additional_kwargs={}, response_metadata={}, id='run-b75d8338-57b6-43b2-8f88-a81389844523')},\n",
       "  'run_id': 'b75d8338-57b6-43b2-8f88-a81389844523',\n",
       "  'name': 'ChatOpenAI',\n",
       "  'tags': ['seq:step:1'],\n",
       "  'metadata': {'ls_provider': 'openai',\n",
       "   'ls_model_name': 'qwen-max',\n",
       "   'ls_model_type': 'chat',\n",
       "   'ls_temperature': 0.7},\n",
       "  'parent_ids': ['1a997380-c991-4732-bb73-7540afcdf5c2']}]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "events = []\n",
    "async for event in chain.astream_events(\"hello\", version=\"v2\"):\n",
    "    events.append(event)\n",
    "\n",
    "events[:3]\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df27199f14823da9",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "嘿，API 中那个奇怪的 version=\"v1\" 参数是什么？！ 😾\n",
    "这是一个beta API，我们几乎肯定会对其进行一些更改。\n",
    "该版本参数将允许我们最小化对您代码的此类更改。\n",
    "\n",
    "如果你检查前几个事件，你会注意到有**3**个不同的开始事件，而不是**2**个开始事件。\n",
    "\n",
    "这三个开始事件对应于：\n",
    "\n",
    "1. 链（模型 + 解析器）\n",
    "2. 模型\n",
    "3. 解析器\n",
    "\n",
    "中间的事件呢？\n",
    "让我们使用这个API来输出模型和解析器的流事件。我们忽略链中的开始事件、结束事件和事件。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "1360e59c75859c3e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T09:41:58.097208Z",
     "start_time": "2024-10-25T09:41:53.296410Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Chat model chunk: ''\n",
      "Chat model chunk: 'Certainly!'\n",
      "Chat model chunk: ' Below'\n",
      "Chat model chunk: ' is'\n",
      "Chat model chunk: ' the information for France'\n",
      "Chat model chunk: ', Spain, and'\n",
      "Chat model chunk: ' Japan, formatted in'\n",
      "Chat model chunk: ' JSON with the structure'\n",
      "Chat model chunk: ' you specified:\\n\\n```'\n",
      "Chat model chunk: 'json\\n{\\n '\n",
      "Parser chunk: {}\n",
      "Chat model chunk: ' \"countries\": [\\n'\n",
      "Parser chunk: {'countries': []}\n",
      "Chat model chunk: '    {\\n      \"name\": \"France'\n",
      "Parser chunk: {'countries': [{'name': 'France'}]}\n",
      "Chat model chunk: '\",\\n      \"population'\n",
      "Chat model chunk: '\": 67'\n",
      "Parser chunk: {'countries': [{'name': 'France', 'population': 67}]}\n",
      "Chat model chunk: '0220'\n",
      "Parser chunk: {'countries': [{'name': 'France', 'population': 670220}]}\n",
      "Chat model chunk: '00\\n   '\n",
      "Parser chunk: {'countries': [{'name': 'France', 'population': 67022000}]}\n",
      "Chat model chunk: ' },\\n    {\\n     '\n",
      "Parser chunk: {'countries': [{'name': 'France', 'population': 67022000}, {}]}\n",
      "Chat model chunk: ' \"name\": \"'\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "num_events = 0\n",
    "async for event in chain.astream_events(\n",
    "        \"output a list of the countries france, spain and japan and their populations in JSON format. Use a dict with an outer key of `countries` which contains a list of countries. Each country should have the key `name` and `population`\",\n",
    "        version=\"v2\"):\n",
    "    kind = event[\"event\"]\n",
    "    if kind == \"on_chat_model_stream\":\n",
    "        print(\n",
    "            f\"Chat model chunk: {repr(event['data']['chunk'].content)}\",\n",
    "            flush=True,\n",
    "        )\n",
    "    if kind == \"on_parser_stream\":\n",
    "        print(\n",
    "            f\"Parser chunk: {event['data']['chunk']}\",\n",
    "            flush=True,\n",
    "        )\n",
    "    num_events += 1\n",
    "\n",
    "    if num_events > 30:\n",
    "        # Break after 30 events\n",
    "        print(\"...\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0f7687a02bc00c3",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 过滤事件\n",
    "由于这个API产生了如此多的事件，能够对事件进行过滤是很有用的。\n",
    "你可以通过组件的name、组件的tags或组件的type进行过滤。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dda9e212d8df883",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 通过名称"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "63d6854f74f9bb4f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T09:46:12.052264Z",
     "start_time": "2024-10-25T09:46:06.614859Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'event': 'on_parser_start', 'data': {'input': 'output a list of the countries france, spain and japan and their populations in JSON format. Use a dict with an outer key of `countries` which contains a list of countries. Each country should have the key `name` and `population`'}, 'name': 'my_parser', 'tags': ['seq:step:2'], 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'metadata': {}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "{'event': 'on_parser_stream', 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'name': 'my_parser', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': {}}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "{'event': 'on_parser_stream', 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'name': 'my_parser', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': {'countries': []}}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "{'event': 'on_parser_stream', 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'name': 'my_parser', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': {'countries': [{}]}}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "{'event': 'on_parser_stream', 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'name': 'my_parser', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France'}]}}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "{'event': 'on_parser_stream', 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'name': 'my_parser', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 670}]}}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "{'event': 'on_parser_stream', 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'name': 'my_parser', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 6702200}]}}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "{'event': 'on_parser_stream', 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'name': 'my_parser', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 67022000}]}}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "{'event': 'on_parser_stream', 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'name': 'my_parser', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 67022000}, {}]}}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "{'event': 'on_parser_stream', 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'name': 'my_parser', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 67022000}, {'name': 'Spain'}]}}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "{'event': 'on_parser_stream', 'run_id': '8b2720dd-5bdf-410b-83da-c4746d88bbf8', 'name': 'my_parser', 'tags': ['seq:step:2'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 67022000}, {'name': 'Spain', 'population': 47}]}}, 'parent_ids': ['439f61f5-1ed8-41e6-b125-fea0ed4c6790']}\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "chain = llm.with_config({\"run_name\": \"my_llm\"}) | JsonOutputParser().with_config({\"run_name\": \"my_parser\"})\n",
    "num_events = 0\n",
    "async for event in chain.astream_events(\n",
    "        \"output a list of the countries france, spain and japan and their populations in JSON format. Use a dict with an outer key of `countries` which contains a list of countries. Each country should have the key `name` and `population`\",\n",
    "        version=\"v2\", include_names=[\"my_parser\"]):\n",
    "    print(event)\n",
    "    num_events += 1\n",
    "    if num_events > 10:\n",
    "        # 截断输出\n",
    "        print(\"...\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b2afbc70d7ab0d",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 通过类型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5018fd58b366d8b1",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "num_events = 0\n",
    "async for event in chain.astream_events(\n",
    "        \"output a list of the countries france, spain and japan and their populations in JSON format. Use a dict with an outer key of `countries` which contains a list of countries. Each country should have the key `name` and `population`\",\n",
    "        version=\"v2\", include_types=[\"chat_model\"]):\n",
    "    print(event)\n",
    "    num_events += 1\n",
    "    if num_events > 10:\n",
    "        # 截断输出\n",
    "        print(\"...\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92bbe20c050cf449",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 通过标签\n",
    "标签会被给定可运行组件的子组件继承。\n",
    "如果你使用标签进行过滤，请确保这是你想要的。 ::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "9d77ab60dbd20f15",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T09:49:56.006142Z",
     "start_time": "2024-10-25T09:49:47.713727Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'event': 'on_parser_start', 'data': {'input': 'output a list of the countries france, spain and japan and their populations in JSON format. Use a dict with an outer key of `countries` which contains a list of countries. Each country should have the key `name` and `population`'}, 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'metadata': {}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "{'event': 'on_parser_stream', 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'metadata': {}, 'data': {'chunk': {}}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "{'event': 'on_parser_stream', 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'metadata': {}, 'data': {'chunk': {'countries': [{}]}}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "{'event': 'on_parser_stream', 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France'}]}}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "{'event': 'on_parser_stream', 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 6741}]}}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "{'event': 'on_parser_stream', 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 67413000}]}}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "{'event': 'on_parser_stream', 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 67413000}, {}]}}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "{'event': 'on_parser_stream', 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 67413000}, {'name': 'Spain'}]}}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "{'event': 'on_parser_stream', 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 67413000}, {'name': 'Spain', 'population': 474}]}}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "{'event': 'on_parser_stream', 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 67413000}, {'name': 'Spain', 'population': 4745079}]}}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "{'event': 'on_parser_stream', 'run_id': '4aa0621e-3502-413d-bd1f-f94b370cbe8c', 'name': 'JsonOutputParser', 'tags': ['seq:step:2', 'my_qa_tag'], 'metadata': {}, 'data': {'chunk': {'countries': [{'name': 'France', 'population': 67413000}, {'name': 'Spain', 'population': 47450795}]}}, 'parent_ids': ['beba172e-f441-4d5b-a0ff-0397a97261eb']}\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "chain = llm.with_config({\"run_name\": \"my_llm\"}) | JsonOutputParser().with_config({\"tags\": [\"my_qa_tag\"]})\n",
    "\n",
    "num_events = 0\n",
    "async for event in chain.astream_events(\n",
    "        \"output a list of the countries france, spain and japan and their populations in JSON format. Use a dict with an outer key of `countries` which contains a list of countries. Each country should have the key `name` and `population`\",\n",
    "        version=\"v2\", include_tags=[\"my_qa_tag\"]):\n",
    "    print(event)\n",
    "    num_events += 1\n",
    "    if num_events > 10:\n",
    "        # 截断输出\n",
    "        print(\"...\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9fbb9ac78d611b0",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## 传播回调\n",
    ":::{.callout-caution} 如果在工具中调用可运行对象，请将回调传播给可运行对象；否则，将不会生成流式事件。 :::\n",
    ":::{.callout-note} 当使用RunnableLambdas或@chain装饰器时，回调会在幕后自动传播。 :::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "3e2ecb243ce062c3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T09:57:58.036675Z",
     "start_time": "2024-10-25T09:57:58.003946Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'event': 'on_tool_start', 'data': {'input': 'hello'}, 'name': 'bad_tool', 'tags': [], 'run_id': 'cbc56500-685a-4828-8967-4f2f6d97f572', 'metadata': {}, 'parent_ids': []}\n",
      "{'event': 'on_chain_start', 'data': {'input': 'hello'}, 'name': 'reverse_word', 'tags': [], 'run_id': '22bba79c-8eea-4fa5-b0bb-915ac444b9ab', 'metadata': {}, 'parent_ids': ['cbc56500-685a-4828-8967-4f2f6d97f572']}\n",
      "{'event': 'on_chain_end', 'data': {'output': 'hello', 'input': 'hello'}, 'run_id': '22bba79c-8eea-4fa5-b0bb-915ac444b9ab', 'name': 'reverse_word', 'tags': [], 'metadata': {}, 'parent_ids': ['cbc56500-685a-4828-8967-4f2f6d97f572']}\n",
      "{'event': 'on_tool_end', 'data': {'output': 'hello'}, 'run_id': 'cbc56500-685a-4828-8967-4f2f6d97f572', 'name': 'bad_tool', 'tags': [], 'metadata': {}, 'parent_ids': []}\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.tools import tool\n",
    "from langchain_core.runnables import RunnableLambda\n",
    "\n",
    "\n",
    "def reverse_word(word: str):\n",
    "    \"\"\"\n",
    "     Reverse the\n",
    "    :param word: \n",
    "    :return: \n",
    "    \"\"\"\n",
    "    return word[::1]\n",
    "\n",
    "reverse_word = RunnableLambda(reverse_word)\n",
    "\n",
    "@tool\n",
    "def bad_tool(word:str):\n",
    "    \"\"\"\n",
    "    不传播回调的自定义工具\n",
    "    :param word: \n",
    "    :return: \n",
    "    \"\"\"\n",
    "    return reverse_word.invoke(word)\n",
    "\n",
    "async for event in bad_tool.astream_events(\"hello\",version=\"v2\"):\n",
    "    print(event)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "492d4bf52a92f2b5",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "这是一个正确传播回调的重新实现。现在你会注意到我们从reverse_word可运行对象中也得到了事件。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "38c718114893a83f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T10:00:41.592750Z",
     "start_time": "2024-10-25T10:00:41.575046Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'event': 'on_tool_start', 'data': {'input': 'hello'}, 'name': 'correct_tool', 'tags': [], 'run_id': '5a840c09-6661-4a85-b382-0fdc4a224a75', 'metadata': {}, 'parent_ids': []}\n",
      "{'event': 'on_chain_start', 'data': {'input': 'hello'}, 'name': 'reverse_word', 'tags': [], 'run_id': '4903b84f-356e-4708-b92c-07cca2528321', 'metadata': {}, 'parent_ids': ['5a840c09-6661-4a85-b382-0fdc4a224a75']}\n",
      "{'event': 'on_chain_end', 'data': {'output': 'hello', 'input': 'hello'}, 'run_id': '4903b84f-356e-4708-b92c-07cca2528321', 'name': 'reverse_word', 'tags': [], 'metadata': {}, 'parent_ids': ['5a840c09-6661-4a85-b382-0fdc4a224a75']}\n",
      "{'event': 'on_tool_end', 'data': {'output': 'hello'}, 'run_id': '5a840c09-6661-4a85-b382-0fdc4a224a75', 'name': 'correct_tool', 'tags': [], 'metadata': {}, 'parent_ids': []}\n"
     ]
    }
   ],
   "source": [
    "@tool\n",
    "def correct_tool(word:str,callbacks):\n",
    "    \"\"\"\n",
    "    正确传播回调的工具。\n",
    "    :param word: \n",
    "    :param callbacks: \n",
    "    :return: \n",
    "    \"\"\"\n",
    "    return reverse_word.invoke(word,{\"callbacks\":callbacks})\n",
    "\n",
    "async for event in correct_tool.astream_events(\"hello\",version=\"v2\"):\n",
    "    print(event)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fded09c8697d94c3",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "如果你从Runnable Lambdas或@chains中调用可运行对象，那么回调将自动传递。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "ae2237c79442e1de",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-25T10:13:25.366798Z",
     "start_time": "2024-10-25T10:13:25.356498Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'event': 'on_chain_start', 'run_id': 'd2cbd064-997b-4401-8461-194a2afc31b5', 'name': 'reverse_and_double', 'tags': [], 'metadata': {}, 'data': {'input': '1234'}, 'parent_ids': []}\n",
      "{'event': 'on_chain_start', 'name': 'reverse_word', 'run_id': '5b00dc26-b116-43d9-922f-d5228360449f', 'tags': [], 'metadata': {}, 'data': {'input': '1234'}, 'parent_ids': []}\n",
      "{'event': 'on_chain_end', 'name': 'reverse_word', 'run_id': '5b00dc26-b116-43d9-922f-d5228360449f', 'tags': [], 'metadata': {}, 'data': {'input': '1234', 'output': '1234'}, 'parent_ids': []}\n",
      "{'event': 'on_chain_stream', 'run_id': 'd2cbd064-997b-4401-8461-194a2afc31b5', 'tags': [], 'metadata': {}, 'name': 'reverse_and_double', 'data': {'chunk': '12341234'}, 'parent_ids': []}\n",
      "{'event': 'on_chain_end', 'name': 'reverse_and_double', 'run_id': 'd2cbd064-997b-4401-8461-194a2afc31b5', 'tags': [], 'metadata': {}, 'data': {'output': '12341234'}, 'parent_ids': []}\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.runnables import RunnableLambda\n",
    "from langchain_core.runnables import chain\n",
    "\n",
    "def reverse_word(word: str):\n",
    "    \"\"\"\n",
    "     Reverse the\n",
    "    :param word: \n",
    "    :return: \n",
    "    \"\"\"\n",
    "    return word[::1]\n",
    "\n",
    "reverse_word = RunnableLambda(reverse_word)\n",
    "\n",
    "@chain\n",
    "async def reverse_and_double(word: str):\n",
    "    return await reverse_word.ainvoke(word) * 2\n",
    "\n",
    "# await reverse_and_double.ainvoke(\"1234\")\n",
    "\n",
    "async for event in reverse_and_double.astream_events(\"1234\", version=\"v1\"):\n",
    "    print(event)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d55139aac58d486d",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
